import numpy as np
import cv2

box_chu=[320,220,400,300]

def xyxy_xywh(box):
    boxed=[None]*4
    # 目标物中心点坐标以及长宽
    boxed[0] = (box[0] + box[2]) / 2
    boxed[1] = (box[1] + box[3]) / 2
    boxed[2] = abs(box[2] - box[0])
    boxed[3] = abs(box[3] - box[1])
    return boxed

def xywh_xyxy(xywh):
    x1 = xywh[0] - xywh[2]//2
    y1 = xywh[1] - xywh[3]//2
    x2 = xywh[0] + xywh[2] // 2
    y2 = xywh[1] + xywh[3] // 2
    return [x1, y1, x2, y2]

def updata_trace_list(box_xin, trace_list, max_list=50):
    if len(trace_list) <= max_list:
        trace_list.append(box_xin)
    else:
        trace_list.pop(0)
        trace_list.append(box_xin)
    return trace_list

class kem:
    def __init__(self,box=box_chu):
        boxed=xyxy_xywh(box)
        k_box=np.array([[boxed[0],boxed[1],boxed[2],boxed[3],0,0]]).T
        #参数初始
        self.A = np.array([[1, 0, 0, 0, 1, 0],
                      [0, 1, 0, 0, 0, 1],
                      [0, 0, 1, 0, 0, 0],
                      [0, 0, 0, 1, 0, 0],
                      [0, 0, 0, 0, 1, 0],
                      [0, 0, 0, 0, 0, 1]])
        self.H=np.eye(6)
        self.Q=np.eye(6)*0.1
        self.R=np.eye(6)*1
        self.P=np.eye(6)
        self.B=None
        #初始卡尔曼数据
        self.X_posterior=np.array(k_box)
        self.P_posterior=np.array(self.P)
        self.Z=np.array(k_box)
        self.guiji=[]

    def draw_trace(self,img,b,g,r):
        for i, item in enumerate(self.guiji):
            if i < 1:
                continue
            cv2.line(img,(self.guiji[i][0], self.guiji[i][1]), (self.guiji[i - 1][0], self.guiji[i - 1][1]),(b,g,r), 3)

    def kem_get_iou(self,qiu):
        #Z
        qiued=xyxy_xywh(qiu)
        #轨迹更新
        box_center = (int(qiued[0]), int(qiued[1]))
        self.guiji = updata_trace_list(box_center, self.guiji, 20)
        self.Z[0:4]=np.array([qiued]).T
        self.Z[4::]=np.array([qiued[0] - self.X_posterior[0],qiued[1]-self.X_posterior[1]])
        # print("Z---------",Z[4::],"\n")
        # -----进行先验估计-----------------
        X_prior = np.dot(self.A, self.X_posterior)
        box_prior = xywh_xyxy(X_prior[0:4])
        # -----计算状态估计协方差矩阵P--------
        P_prior_1 = np.dot(self.A, self.P_posterior)
        P_prior = np.dot(P_prior_1, self.A.T) + self.Q
        # ------计算卡尔曼增益---------------------
        k1 = np.dot(P_prior, self.H.T)
        k2 = np.dot(np.dot(self.H, P_prior), self.H.T) + self.R
        K = np.dot(k1, np.linalg.inv(k2))
        # --------------后验估计------------
        X_posterior_1 = self.Z - np.dot(self.H, X_prior)
        self.X_posterior = X_prior + np.dot(K, X_posterior_1)
        box_posterior = xywh_xyxy(self.X_posterior[0:4])
        # ---------更新状态估计协方差矩阵P-----
        P_posterior_1 = np.eye(6) - np.dot(K, self.H)
        self.P_posterior = np.dot(P_posterior_1, P_prior)
        return box_posterior

    def kem_no_iou(self):
        # 如果IOU匹配失败，此时失去观测值，那么直接使用上一次的最优估计作为先验估计
        # 此时直接迭代，不使用卡尔曼滤波
        self.X_posterior = np.dot(self.A, self.X_posterior)
        box_posterior = xywh_xyxy(self.X_posterior[0:4])
        box_center = (
            (int(box_posterior[0] + box_posterior[2]) // 2), int((box_posterior[1] + box_posterior[3]) // 2))
        self.guiji = updata_trace_list(box_center, self.guiji, 20)
        return box_posterior
